import numpy as np
from pyE17 import utils as U
from pyE17 import io
import ptycho
import os

# Main parameters
energy = 6.2
z = 7.046        # Distance chosen to have exactly 32 nm pixel size
ds = 172e-6
lam = 1.2398e-9/energy

# Number of incident photons in the probe
noise = 3e6

# Carbon, thick
ior = -1j*3.4054267E-08 - 1.19337546E-05
max_thickness = 1e-6

run_name = 'lownoise_no_displacements'

subpix_method = 'linear'

# Displacement rms in pixels
disp_rms = 0. 

# subpix functions are not exposed in ptycho (maybe this should be done)
def subpix_fourier(a,ramp,forward=True):
    """\
    Perform a shift with the provided ramp.
    """
    if forward:
        return np.fft.ifftn(np.fft.fftn(a)*ramp)
    else:
        return np.fft.ifftn(np.fft.fftn(a)*ramp.conj())

def subpix_linear(a, pts, forward=True):
    """\
    Perform a subpixel shift using bilinear interpolation.
    """
    x0,x1 = pts
    if not forward: x0,x1 = -x0,-x1
    a00 = U.pshift(a,[np.floor(x0),np.floor(x1)],'nearest')
    a01 = U.pshift(a,[np.floor(x0),np.ceil(x1)],'nearest')
    a10 = U.pshift(a,[np.ceil(x0),np.floor(x1)],'nearest')
    a11 = U.pshift(a,[np.ceil(x0),np.ceil(x1)],'nearest')
    x0 = x0 % 1
    x1 = x1 % 1
    return a00*(1-x0)*(1-x1) + a01*(1-x0)*x1 + a10*x0*(1-x1) + a11*x0*x1


probe_file = 'probe_1_128.h5'
object_file = 'object_biosample.h5'

probe = io.h5read(probe_file)['probe']
probe /= U.norm(probe)
ob0 = io.h5read(object_file)['object']
object = np.exp(-2j * np.pi * ior * max_thickness * ob0 / lam)

asize = probe.shape
dx_spec = lam*z/(asize[0]*ds) 

nth = 5
dr = .7e-6
lx = 12e-6
ly = 12e-6

positions = ptycho.round_scan_ROI_positions(dr, lx, ly, nth)

positions = np.array(positions) / dx_spec
positions -= positions.min(axis=0)
positions_float = positions
positions = np.round(positions_float)
Npos = len(positions)

# Create the random displacements
displacements = disp_rms * np.random.normal(size=positions.shape)

# Here we can combine the displacements in subpix_parameters
if subpix_method.lower() == 'fourier':
    p_shift = positions - positions_float - displacements
    subpix_parameter = [np.exp(2j * np.pi * np.sum(U.fgrid(asize,x/asize),axis=0)) for x in p_shift]
elif subpix_method.lower() == 'linear':
    subpix_parameter = positions - positions_float - displacements

# This is the way the final object size is computed 
object_size = tuple(np.array(asize) + positions.max(axis=0))

# Setup object views for computation of diffraction patterns
offset = (1000,1550)
obj = object[offset[0]:(offset[0] + object_size[0]), offset[1]:(offset[1] + object_size[1])]
obview = [obj[p[0]:(p[0] + asize[0]), p[1]:(p[1] + asize[1])] for p in positions]
coverage = np.zeros_like(obj).astype(float)
for p in positions:
    coverage[p[0]:(p[0] + asize[0]), p[1]:(p[1] + asize[1])] += abs(probe)**2

io.h5write('simul_%s_solution.h5' % run_name, obj = obj.copy(), probe = probe, coverage=coverage)

Iprobe = (np.abs(np.fft.fftn(probe))**2).sum()

data = np.zeros((Npos,) + asize, dtype=int)
for ii,iob in enumerate(range(len(obview))):
    if subpix_method == 'fourier':
        pr = subpix_fourier(probe, subpix_parameter[ii])
    elif subpix_method == 'linear':
        pr = subpix_linear(probe, subpix_parameter[ii])
    dp = np.fft.fftshift(np.fft.fftn(pr * obview[iob]))
    data[iob,:,:] = np.random.poisson(np.abs(dp)**2 * noise / Iprobe)
dirname = 'simul_%s' % run_name
try:
    os.mkdir(dirname)
except OSError:
    pass
filename = dirname + ('/simul_%s_data_%dx%d.h5' % (run_name, asize[0], asize[1]))
print filename
io.h5write(filename, data=data, fmask=np.ones(asize,dtype=int), displacements=displacements)
